import Data.Char (isDigit)
import qualified Data.Map as Map
import Data.Map (Map)
import Utils (tok, dropWhileIncl)

type Part    = [Int]
type PartIdx = Int

type WorkflowMap  = Map WorkflowName Workflow
type WorkflowName = String
type Workflow     = [Either Condition WorkflowName]
type Condition    = (PartIdx, Operation, Operand, WorkflowName)
data Operation    = Less | Greater deriving (Show, Eq)
type Operand      = Int

parseInput :: [String] -> (WorkflowMap, [Part])
parseInput ls = 
    (wfm, parts)
    where
        wfm    = Map.fromList . map parseWorkflow $ blocks !! 0
        parts  = map parsePart $ blocks !! 1
        blocks = tok [""] ls

-- px{a<2006:qkq,m>2090:A,rfg}
parseWorkflow :: String -> (WorkflowName, Workflow)
parseWorkflow inp = 
    (name, steps)
    where
        name = takeWhile (/='{') $ inp
        steps = map parseStep . tok "," . init . dropWhileIncl (/='{') $ inp
            where
                parseStep s
                    | '<' `elem` s = Left (partIdx s, Less, operand s, goto s)
                    | '>' `elem` s = Left (partIdx s, Greater, operand s, goto s)
                    | otherwise    = Right s
                partIdx s = case (head s) of {'x' -> 0; 'm' -> 1; 'a' -> 2; 's' -> 3}
                operand   = read . takeWhile (isDigit) . dropWhileIncl (`notElem` "><")
                goto      = takeWhile (/=',') . dropWhileIncl (/=':')

parsePart :: String -> Part
parsePart inp = numbers 
    where
        toks    = map read . tok "{},xmas=" $ inp
        numbers = [toks!!0, toks!!1, toks!!2, toks!!3]

decide :: WorkflowMap -> Workflow -> Part -> Bool
decide _ ((Right "A"):_) = True
decide _ ((Right "R"):_) = False
decide wfm ((Right s):_) pt = decide wfm (wfm Map.! s) pt
decide wfm ((Left (idx, op, opd, goto)):steps) pt
    | op == Less && pt !! idx < opd = decide wfm [(Right goto)]
    | op == Greater && pt !! idx > opd = decide wfm [(Right goto)]
    | otherwise = decide wfm steps pt

main = do
    contents <- readFile "input.txt"
    let (workflowMap, parts) = parseInput . lines $ contents
    let firstworkflow = workflowMap Map.! "in"

    print $ sum . map sum . filter (decide workflowMap firstworkflow) $ parts
